# https://gitee.com/yueyinqiu5990/tj12413601/blob/master/assignment2/question2/main_spatial.py
import math

import matplotlib.pyplot
import numpy
import scipy

import hydrogen
import my_special


def _calculate_and_draw(subplot, title, hydrogen_atom: hydrogen.Hydrogen):
    r = numpy.linspace(0, 1e-9, 1000)
    theta = numpy.linspace(0, 2 * math.pi, 1000)
    r, theta = numpy.meshgrid(r, theta)

    radial = hydrogen_atom.radial(r)
    p = hydrogen_atom.poloidal(0, theta)

    wave = numpy.abs(radial * p) ** 2

    x = r * numpy.sin(theta)
    y = r * numpy.cos(theta)

    matplotlib.pyplot.subplot(*subplot)
    matplotlib.pyplot.title(title)
    matplotlib.pyplot.axis("off")
    matplotlib.pyplot.contourf(x, y, wave, levels=100)


def main():
    h = hydrogen.Hydrogen(
        assoc_laguerre_function=my_special.assoc_laguerre,
        sph_harm_function=my_special.sph_harm
    )
    for m in range(0, 3):
        h.set_state(5, 2, m)
        _calculate_and_draw((2, 3, m + 1), f"n=5, l=2, m={m}", h)

    h = hydrogen.Hydrogen(
        assoc_laguerre_function=scipy.special.assoc_laguerre,
        sph_harm_function=scipy.special.sph_harm
    )
    for m in range(0, 3):
        h.set_state(5, 2, m)
        _calculate_and_draw((2, 3, m + 1 + 3), "With SciPy", h)

    matplotlib.pyplot.subplots_adjust(hspace=0.25, wspace=0.25, right=0.717)
    matplotlib.pyplot.show()


if __name__ == "__main__":
    main()
